#include "simulation.h"

// Weight normalization: makes largest-valued weight 1.0 or the smallest-valued weight -1.0
// and scales other weights accordingly
__host__ __device__
void normalize_weights(const int num_weights, float *weights)
{
  // 1) find the largest element by absolute value
  unsigned int index_of_largest = 0;
  for(unsigned int i = 1; i < num_weights; i++) // iterate from 1 since first test is between 0 and 1
  {
    if(abs(weights[i]) > abs(weights[index_of_largest]))
    {
      index_of_largest = i;
    }
  }
  
  // 2) divide the whole array by that value
  float largest = abs(weights[index_of_largest]);
  if(largest == 0.0)
    return; // never divide by zero
  for(unsigned int i = 0; i < num_weights; i++)
  {
    weights[i] /= largest;
  }
}

// kernel normalizes the weights of a single creature
__global__
void normalization_kernel(const int num_weights,
                          const int num_creatures,
                          float *creature_weights)
{
  // alloc num_weights * sizeof(float) * block_size shared memory in kernel launch
  extern __shared__ float smem[];
  unsigned int i = blockDim.x * blockIdx.x + threadIdx.x;
  const unsigned int grid_size = gridDim.x * blockDim.x;
  float *s_weights = &smem[threadIdx.x * num_weights];

  for(; i < num_creatures; i += grid_size)
  {
    // copy from global to shared mem
    float *g_weights = &creature_weights[i * num_weights];
    for(int j=0; j < num_weights; j++)
    {
      s_weights[j] = g_weights[j];
    }

    // perform the normalization in shared mem
    normalize_weights(num_weights, s_weights);

    // write back out from shared mem to global
    for(int j=0; j < num_weights; j++)
    {
      g_weights[j] = s_weights[j];
    }
  }
}

void normalize_all_weights(const int num_creatures,
                           const int num_weights,
                           thrust::device_vector<float> &d_creature_weights,
                           thrust::host_vector<float> &h_creature_weights)
{
    unsigned int block_size = 512;
    unsigned int grid_size = (num_creatures + block_size - 1) / block_size;
    normalization_kernel<<<grid_size, block_size, num_weights * sizeof(float) * block_size>>>(num_weights, num_creatures, thrust::raw_pointer_cast(&d_creature_weights[0]));
    h_creature_weights = d_creature_weights;
}

// kernel interacts two agents as part of a round-robin tournament among all agents
__global__
void simulation_kernel(const int memory_length,
                       const int num_weights,
                       const int num_creatures,
                       const int num_rounds_per_match,
                       unsigned int *random,
                       const float *creature_weights,
                       int *creature_scores,
                       int R,
                       int S,
                       int T,
                       int P,
                       const bool use_const,
                       const bool use_random,
                       const bool use_same)
{
  extern __shared__ float s_weights[];
  unsigned int i = blockDim.x * blockIdx.x + threadIdx.x;
  const unsigned int grid_size = gridDim.x * blockDim.x;
  unsigned int memory;
  
  // load the weights into shared memory
  for(unsigned int j=threadIdx.x; j < num_creatures * num_weights; j += blockDim.x)
  {
    s_weights[j] = creature_weights[j];
  }
  
  __syncthreads();

  unsigned int num_pairs = utlm_num_elements(num_creatures);
  for(; i < num_pairs; i += grid_size)
  {
    reset_memory(memory);

    unsigned int ca = utlm_index_to_row(i, num_creatures);
    unsigned int cb = utlm_index_to_col(i, ca, num_creatures);

    unsigned int ra, rb;
    if(use_random) {
      ra = random[ca*num_creatures+cb];
      rb = random[cb*num_creatures+ca];
    }

    int scorea = 0;
    int scoreb = 0;
    for(int r = 0; r < num_rounds_per_match; r++) {
      // calculate the decisons: cooperate or defect
      float asum = 0.0f;
      float bsum = 0.0f;
      unsigned int m;
      float wa, wb;
      for(m = 0; m < memory_length; m++)
      {
        unsigned int mem = get_memory(memory, m);
        wa = s_weights[ca * num_weights + m];
        wb = s_weights[cb * num_weights + m];
        asum += wa * mem_to_b(mem);
        bsum += wb * mem_to_a(mem);
      }

      // sigmoid transform
      asum = 2.0f / (1.0f + exp(-asum)) - 1.0f;
      bsum = 2.0f / (1.0f + exp(-bsum)) - 1.0f;

      if (use_const)
      {
        wa = s_weights[ca * num_weights + m];
        wb = s_weights[cb * num_weights + m];
        asum += wa;
        bsum += wb;
        m++;
      }

      if (use_random)
      {
        ra = next_rand(ra);
        rb = next_rand(rb);
        wa = s_weights[ca * num_weights + m] * (float)ra / (float)max_rand();
        wb = s_weights[cb * num_weights + m] * (float)rb / (float)max_rand();
        asum += wa;
        bsum += wb;
        m++;
      }

      if (use_same)
      {
        float dot = 0.0f;
        for(unsigned int w = 0; w < num_weights; w++)
        {
          dot += s_weights[ca * num_weights + w] * s_weights[cb * num_weights + w];
        }
        asum += s_weights[ca * num_weights + m] * dot;
        bsum += s_weights[cb * num_weights + m] * dot;
        m++;
      }

      // determine the scores
      int d = ab_to_mem(asum, bsum);
      switch(d)
      {
        case 1: scorea += R; scoreb += R; break;
        case 2: scorea += S; scoreb += T; break;
        case 3: scorea += T; scoreb += S; break;
        case 4: scorea += P; scoreb += P; break;
      }

      push_memory(memory, d);
    }

    creature_scores[ca*num_creatures+cb] = scorea;
    creature_scores[cb*num_creatures+ca] = scoreb; // hack  (ca==cb) ? scorea + scoreb :

    if(use_random) {
      random[ca*num_creatures+cb] = ra;
      random[cb*num_creatures+ca] = rb;
    }
  }
}


void host_simulation(const int memory_length,
                     const int num_weights,
                     const int num_creatures,
                     const int num_rounds_per_match,
                     thrust::host_vector<unsigned int> &h_random,
                     const thrust::host_vector<float> &h_creature_weights,
                     thrust::host_vector<int> &h_creature_scores,
                     int R,
                     int S,
                     int T,
                     int P,
                     const bool use_const,
                     const bool use_random,
                     const bool use_same)
{
  thrust::host_vector<int> h_memory(memory_length);

  for(unsigned int ca = 0; ca < num_creatures; ca++) {
    for(unsigned int cb = ca + 1; cb < num_creatures; cb++) {
      thrust::fill(h_memory.begin(), h_memory.end(), (int) 0);

      unsigned int ra, rb;
      if(use_random) {
        ra = h_random[ca*num_creatures+cb];
        rb = h_random[cb*num_creatures+ca];
      }

      for(int r = 0; r < num_rounds_per_match; r++) {
        float asum = 0.0f;
        float bsum = 0.0f;
        unsigned int m;
        float wa, wb;
        for(m = 0; m < memory_length; m++)
        {
          wa = h_creature_weights[ca * num_weights + m];
          wb = h_creature_weights[cb * num_weights + m];
          asum += wa * mem_to_b(h_memory[m]);
          bsum += wb * mem_to_a(h_memory[m]);
//          std::cout << "m=" << m << " wa=" << wa << " wb=" << wb << " asum=" << asum << " bsum=" << bsum << " memory=" << h_memory[m] << std::endl;
        }

        // sigmoid transform
        asum = 2.0f / (1.0f + exp(-asum)) - 1.0f;
        bsum = 2.0f / (1.0f + exp(-bsum)) - 1.0f;

        if (use_const)
        {
          wa = h_creature_weights[ca * num_weights + m];
          wb = h_creature_weights[cb * num_weights + m];
          asum += wa;
          bsum += wb;
//          std::cout << "uc " << " wa=" << wa << " wb=" << wb << " asum=" << asum << " bsum=" << bsum << std::endl;
          m++;
        }

        if (use_random)
        {
          ra = next_rand(ra);
          rb = next_rand(rb);
//          std::cerr << ra << " " << rb << std::endl;
          wa = h_creature_weights[ca * num_weights + m] * (float)ra / (float)max_rand();
          wb = h_creature_weights[cb * num_weights + m] * (float)rb / (float)max_rand();
          asum += wa;
          bsum += wb;
          m++;
        }

        if (use_same)
        {
          float dot = 0.0f;
          for(unsigned int w = 0; w < num_weights; w++)
          {
            dot += h_creature_weights[ca * num_weights + w] * h_creature_weights[cb * num_weights + w];
          }
          asum += h_creature_weights[ca * num_weights + m] * dot;
          bsum += h_creature_weights[cb * num_weights + m] * dot;
          m++;
        }

        int d = ab_to_mem(asum, bsum);
        int scorea = 0, scoreb = 0;
        switch(d)
        {
          case 1: scorea += R; scoreb += R; break;
          case 2: scorea += S; scoreb += T; break;
          case 3: scorea += T; scoreb += S; break;
          case 4: scorea += P; scoreb += P; break;
        }
        h_creature_scores[ca] += scorea;
        h_creature_scores[cb] += scoreb;

        for(m = memory_length-1; m > 0; m--)
        {
          h_memory[m] = h_memory[m-1];
        }
        h_memory[0] = d;
//           std::cout << "ca=" << ca << " cb=" << cb << " d=" << d << " asum=" << asum << " bsum=" << bsum << std::endl;
      }

      if(use_random) {
        h_random[ca*num_creatures+cb] = ra;
        h_random[cb*num_creatures+ca] = rb;
      }
    }
  }
}

void device_simulation(const int memory_length,
                       const int num_weights,
                       const int num_creatures,
                       const int num_rounds_per_match,
                       thrust::device_vector<unsigned int> &d_random,
                       const thrust::device_vector<float> &d_creature_weights,
                       thrust::device_vector<int> &d_creature_scores,
                       int R,
                       int S,
                       int T,
                       int P,
                       const bool use_const,
                       const bool use_random,
                       const bool use_same)
{
//    unsigned int max_block = (((16 * 1024 / sizeof(float) / 2) / num_weights) >> 5) << 5;
//    block_size = (max_block > 512) ? 512 : max_block;
  static thrust::device_vector<int> d_creature_scores_unreduced(num_creatures*num_creatures);
  static thrust::device_vector<unsigned int> d_reduced_keys(num_creatures); // seems necessary to make the reduction work

  unsigned int block_size = 512;
  unsigned int grid_size = (utlm_num_elements(num_creatures) + block_size - 1) / block_size;
  simulation_kernel<<<grid_size, block_size, num_weights*num_creatures*sizeof(float)>>>(
      memory_length,
      num_weights,
      num_creatures,
      num_rounds_per_match,
      thrust::raw_pointer_cast(&d_random[0]),
      thrust::raw_pointer_cast(&d_creature_weights[0]),
      thrust::raw_pointer_cast(&d_creature_scores_unreduced[0]),
      R,
      S,
      T,
      P,
      use_const,
      use_random,
      use_same
  );

//    std::cout << "Block size: " << block_size << std::endl;

  // report errors
  check_cuda_error("simulation_kernel", __FILE__, __LINE__);
  
  // reduce the individual interaction results
  equal_divide_n equal_divide_n_pred(num_creatures);
  thrust::reduce_by_key(
      thrust::make_counting_iterator(0),
      thrust::make_counting_iterator(num_creatures*num_creatures),
      d_creature_scores_unreduced.begin(),
      d_reduced_keys.begin(),
      d_creature_scores.begin(),
      equal_divide_n_pred
  );
}
